package com.neusoft.databus.agent.interactive.collect;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Vector;
import java.util.stream.Collectors;

import org.snmp4j.PDU;
import org.snmp4j.ScopedPDU;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.mp.SnmpConstants;
import org.snmp4j.smi.OID;
import org.snmp4j.smi.Variable;
import org.snmp4j.smi.VariableBinding;
import org.snmp4j.util.DefaultPDUFactory;
import org.snmp4j.util.PDUFactory;
import org.snmp4j.util.TreeEvent;
import org.snmp4j.util.TreeUtils;

import com.google.common.collect.Multimap;
import com.neusoft.bizcore.web.utils.SpringUtil;
import com.neusoft.databus.agent.parser.SnmpParser;
import com.neusoft.databus.agent.pool.snmp.SnmpConnector;
import com.neusoft.databus.agent.pool.snmp.SnmpWrapper;
import com.neusoft.databus.common.bean.CommandBean;
import com.neusoft.databus.common.bean.MetricBean;
import com.neusoft.databus.common.connection.SnmpConnection;
import com.neusoft.databus.common.exception.DisconnectException;
import com.neusoft.databus.common.parser.Parser;

import lombok.extern.slf4j.Slf4j;

/**
 * snmp collector through snmp
 *
 * @author sunchf
 * @date 2019年1月4日 下午5:31:26
 */
@Slf4j
public class SnmpSnmpCollector extends BasicCollector<SnmpWrapper> {

    public SnmpSnmpCollector(final Parser parser) {
        super(parser);
        this.connector = SpringUtil.getBean(SnmpConnector.class);
    }

    @Override
    public void sendCommand(final CommandBean command) {

        if (SnmpSnmpCollector.log.isDebugEnabled()) {
            SnmpSnmpCollector.log.debug("begin to communicate with snmp agent，{} metrics will be executed",
                    command.getResourceBean().stream().mapToInt(it -> it.getMetrics().size()).sum());
        }

        int batchSize = ((SnmpConnection) command.getConnection()).getBatchSize();
        if (batchSize < 1) {
            batchSize = 100;
        }

        final Multimap<String, MetricBean> multimap = ((SnmpParser) this.parser).parseCollectRequest(command);

        if (SnmpSnmpCollector.log.isDebugEnabled()) {
            SnmpSnmpCollector.log.debug("after optimizing, it will communicate with snmp agent {} times",
                    multimap.keySet().size());
        }

        final Map<String, Variable> fetchedOidValueMap = new HashMap<>();
        final List<VariableBinding> getVariables = multimap.keySet().stream().filter(it -> {
            final MetricBean any = multimap.get(it).stream().findAny().get();
            return (null == any.getRule()) || !any.getRule().containsKey("bulk");
        }).map(c -> {
            return new VariableBinding(new OID(c));
        }).collect(Collectors.toList());

        fetchedOidValueMap.putAll(this.getRequest(this.conn, getVariables, batchSize));

        for (final Entry<String, Variable> entry : fetchedOidValueMap.entrySet()) {
            for (final MetricBean metric : multimap.get(entry.getKey())) {
                metric.setOriginalValue(entry.getValue());
            }
        }

        fetchedOidValueMap.clear();
        final List<String> bulkRoots = multimap.keySet().stream().filter(it -> {
            final MetricBean any = multimap.get(it).stream().findAny().get();
            return (null != any.getRule()) && any.getRule().containsKey("bulk");
        }).collect(Collectors.toList());

        if (bulkRoots.size() > 0) {
            fetchedOidValueMap.putAll(this.getBulkRequest(this.conn, bulkRoots));
            for (final String root : bulkRoots) {
                for (final MetricBean metric : multimap.get(root)) {
                    final String reqCommand = SnmpParser.formatOidStr(metric.getCommand());
                    final Optional<Entry<String, Variable>> opt = fetchedOidValueMap.entrySet().stream()
                            .filter(it -> it.getKey().equals(reqCommand)).findFirst();
                    if (opt.isPresent()) {
                        metric.setOriginalValue(opt.get().getValue());
                    }
                }
            }
        }

        if (SnmpSnmpCollector.log.isDebugEnabled()) {
            SnmpSnmpCollector.log.debug("snmp get all results: {}", fetchedOidValueMap.size());
        }

    }

    /**
     * SNMP GET 操作
     *
     * @param wrapper
     * @param oids
     * @return
     * @throws IOException higher layer will catch IOException
     */
    private Map<String, Variable> getRequest(final SnmpWrapper wrapper, final List<VariableBinding> oids,
            final int batchSize) {

        final Map<String, Variable> result = new HashMap<>();
        final int groupCount =
                (oids.size() % batchSize) == 0 ? (oids.size() / batchSize) : ((oids.size() / batchSize) + 1);

        for (int i = 0; i < groupCount; i++) {
            final List<VariableBinding> subReqList = oids.subList(i * batchSize,
                    ((i + 1) * batchSize) > oids.size() ? oids.size() : (i + 1) * batchSize);
            // 一次批量发送的PDU
            PDU pdu = null;
            if (SnmpConstants.version3 == wrapper.getVersion()) {
                pdu = new ScopedPDU();
            } else {
                pdu = new PDU();
            }

            for (int j = 0; j < subReqList.size(); j++) {
                if (SnmpSnmpCollector.log.isDebugEnabled()) {
                    SnmpSnmpCollector.log.debug("oid: ", subReqList.get(j));
                }
                pdu.add(subReqList.get(j));
            }

            ResponseEvent responseEvent = null;
            try {
                responseEvent = wrapper.get(pdu);
            } catch (final IOException e) {
                // 这里可以判断中断异常
                throw new DisconnectException("the PDU cannot be sent to the target", e);
            }

            if ((responseEvent == null) || (responseEvent.getResponse() == null)) {
                SnmpSnmpCollector.log.warn("snmp get without response");
                continue;
            }

            final Vector<? extends VariableBinding> fetchedVariableBindings = responseEvent.getResponse()
                    .getVariableBindings();

            for (final VariableBinding vb : fetchedVariableBindings) {
                result.put(vb.getOid().toString(), vb.getVariable());
            }
        }

        if (SnmpSnmpCollector.log.isDebugEnabled()) {
            SnmpSnmpCollector.log.debug("snmp get results: {}", result.size());
        }

        return result;
    }

    /**
     * @param wrapper
     * @param roots
     * @return
     */
    private Map<String, Variable> getBulkRequest(final SnmpWrapper wrapper, final List<String> roots) {

        final Map<String, Variable> result = new HashMap<>();
        PDUFactory pduFactory = null;
        if (SnmpConstants.version1 == wrapper.getTarget().getVersion()) {
            pduFactory = new DefaultPDUFactory(PDU.GETNEXT);
        } else {
            pduFactory = new DefaultPDUFactory(PDU.GETBULK);
        }
        final TreeUtils treeUtils = new TreeUtils(wrapper.getSnmp(), pduFactory);

        for (final String root : roots) {
            final List<TreeEvent> treeEvents = treeUtils.getSubtree(wrapper.getTarget(), new OID(root));
            for (final TreeEvent treeEvent : treeEvents) {
                if (null != treeEvent.getVariableBindings()) {
                    for (final VariableBinding vb : treeEvent.getVariableBindings()) {
                        result.put(vb.getOid().toString(), vb.getVariable());
                    }
                }
            }
        }

        return result;
    }

    public static void main(final String[] args) {
        final List<Integer> oids = Arrays.asList(1, 2, 3, 4, 5, 6, 7);
        final int batchSize = 1;
        final int groupCount =
                (oids.size() % batchSize) == 0 ? (oids.size() / batchSize) : ((oids.size() / batchSize) + 1);

        for (int i = 0; i < groupCount; i++) {
            final List<Integer> subReqList = oids.subList(i * batchSize,
                    ((i + 1) * batchSize) > oids.size() ? oids.size() - 1 : (i + 1) * batchSize);
            System.out.println(subReqList);
        }
    }
}
